#include "RHIBuffer.h"
#include "../../Core/CoreGlobals.h"
#include "../../Utils/LogUtil.h"
#include "../RenderState.h"

RHIBuffer::~RHIBuffer()
{
	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		// vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), mBuffers[i], nullptr);
		// vkFreeMemory(GVulkanRHI->GetLogicalDevice(), mBufferMemorys[i], nullptr);
		GRenderState->DelayGCBuffer(mBuffers[i]);
		GRenderState->DelayGCDeviceMemory(mBufferMemorys[i]);
	}
}

VkBufferCreateFlags RHIBuffer::GetBufferCreateFlags() const
{
	return mFlags;
}

VkDeviceSize RHIBuffer::GetSize() const
{
	return mSize;
}

VkBufferUsageFlags RHIBuffer::GetUsage() const
{
	return mUsage;
}

VkSharingMode RHIBuffer::GetSharingMode() const
{
	return mSharingMode;
}

std::vector<VkQueueFlagBits> RHIBuffer::GetQueueFamilys() const
{
	return mQueueFamilys;
}

bool RHIBuffer::IsInQueueFamily(VkQueueFlagBits queueFlag) const
{
	for (const auto queueFamily : mQueueFamilys)
	{
		if (queueFlag == queueFamily)
			return true;
	}
	return false;
}

VkMemoryPropertyFlags RHIBuffer::GetMemoryProperty() const
{
	return mMemoryProperty;
}

bool RHIBuffer::IsMultiple() const
{
	return mIsMultiple;
}

VkBuffer RHIBuffer::GetBuffer(uint32 index /*= UINT32_MAX*/) const
{
	if (!mIsMultiple)
		return mBuffers[0];
	if (index == UINT32_MAX)
		return mBuffers[GetFrameIndex()];
	assert(index <= GetFrameCount());
	return mBuffers[index];
}

VkDeviceMemory RHIBuffer::GetBufferMemory(uint32 index /*= UINT32_MAX*/) const
{
	if (!mIsMultiple)
		return mBufferMemorys[0];
	if (index == UINT32_MAX)
		return mBufferMemorys[GetFrameIndex()];
	assert(index <= GetFrameCount());
	return mBufferMemorys[index];
}

void RHIBuffer::WriteAllMemory(void* buffer, VkDeviceSize offset /*= 0*/, VkDeviceSize size /*= VK_WHOLE_SIZE*/)
{
	assert(mIsMultiple);
	for (uint32 i = 0; i < GetFrameCount(); ++i)
	{
		WriteMemory(buffer, offset, mSize, i);
	}
}

void RHIBuffer::WriteMemory(void* buffer, VkDeviceSize offset, VkDeviceSize size, uint32 index)
{
	// Clac the real size to ceate staging buffer
	size = size == VK_WHOLE_SIZE ? mSize - offset : size;

	// Only propertyFlags is VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT we can vkMapMemory
	// If propertyFlags is VK_MEMORY_PROPERTY_HOST_COHERENT_BIT we need to use vkFlushMappedMemoryRanges and vkInvalidateMappedMemoryRanges to write data
	if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT)
	{
		VkDeviceMemory bufferMemory = GetBufferMemory(index);

		if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
		{
			void* data;
			VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory, offset, size, 0, &data));
			memcpy(data, buffer, size);
			vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory);
		}
		else
		{
			void* data;
			VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory, offset, size, 0, &data));
			memcpy(data, buffer, size);

			// Flush to make changes visible to the host
			VkMappedMemoryRange mappedMemoryRange;
			mappedMemoryRange.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE;
			mappedMemoryRange.memory = bufferMemory;
			mappedMemoryRange.offset = offset;
			mappedMemoryRange.size = size;
			vkFlushMappedMemoryRanges(GVulkanRHI->GetLogicalDevice(), 1, &mappedMemoryRange);
			vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory);
		}
	}
	else if (mMemoryProperty & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
	{
		// We can only use stage buffer to copy data on device
		// Usage need to set VK_BUFFER_USAGE_TRANSFER_DST_BIT
		assert(mUsage & VK_BUFFER_USAGE_TRANSFER_DST_BIT);

		VkBuffer stagingBuffer;
		VkDeviceMemory stagingMemory;

		GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, size, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

		void* data;
		VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, size, 0, &data));
		memcpy(data, buffer, size);
		vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);

		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();


		VkBufferCopy bufferCopy{};
		// Copy whole staging buffer
		bufferCopy.srcOffset = 0;
		bufferCopy.dstOffset = offset;
		bufferCopy.size = size;
		vkCmdCopyBuffer(commandBuffer, stagingBuffer, GetBuffer(index), 1, &bufferCopy);


		GRenderState->FlushSingleCommand(commandBuffer);

		vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
		vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
	}
	else
	{
		//TODO: more ways to copy data
		assert(0);
	}
}

void RHIBuffer::ReadMemory(void* buffer, VkDeviceSize offset, VkDeviceSize size, uint32 index)
{
	// Clac the real size
	size = size == VK_WHOLE_SIZE ? mSize - offset : size;

	if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT)
	{
		VkDeviceMemory bufferMemory = GetBufferMemory(index);

		if (mMemoryProperty & VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
		{
			void* data;
			VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory, offset, size, 0, &data));
			memcpy(buffer, data, size);
			vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory);
		}
		else
		{
			void* data;
			VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory, offset, size, 0, &data));

			// Flush to make changes visible to the host
			VkMappedMemoryRange mappedMemoryRange;
			mappedMemoryRange.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE;
			mappedMemoryRange.memory = bufferMemory;
			mappedMemoryRange.offset = offset;
			mappedMemoryRange.size = size;
			vkInvalidateMappedMemoryRanges(GVulkanRHI->GetLogicalDevice(), 1, &mappedMemoryRange);
			memcpy(buffer, data, size);

			vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), bufferMemory);
		}
	}
	else if (mMemoryProperty & VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
	{
		// Usage need to set VK_BUFFER_USAGE_TRANSFER_DST_BIT to copy data to stagingBuffer
		assert(mUsage & VK_BUFFER_USAGE_TRANSFER_SRC_BIT);

		VkBuffer stagingBuffer;
		VkDeviceMemory stagingMemory;

		GVulkanRHI->CreateBuffer(stagingBuffer, stagingMemory, size, VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);

		VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

		VkBufferCopy bufferCopy{};
		// Copy whole staging buffer
		bufferCopy.srcOffset = 0;
		bufferCopy.dstOffset = offset;
		bufferCopy.size = size;
		vkCmdCopyBuffer(commandBuffer, GetBuffer(index), stagingBuffer, 1, &bufferCopy);

		GRenderState->FlushSingleCommand(commandBuffer);


		void* data;
		VK_CHECK_RESULT(vkMapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, 0, size, 0, &data));
		memcpy(buffer, data, size);
		vkUnmapMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory);



		vkDestroyBuffer(GVulkanRHI->GetLogicalDevice(), stagingBuffer, nullptr);
		vkFreeMemory(GVulkanRHI->GetLogicalDevice(), stagingMemory, nullptr);
	}
	else
	{
		//TODO: more ways to read data
		assert(0);
	}
}

VkDeviceAddress RHIBuffer::GetDeviceAddress() const
{
	assert(mUsage & VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT);

	VkBufferDeviceAddressInfoKHR bufferDeviceAddressInfo{};
	bufferDeviceAddressInfo.sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO;
	bufferDeviceAddressInfo.buffer = mBuffers[GetFrameIndex()];
	VkDeviceAddress deviceAddress = vkGetBufferDeviceAddressKHR(GVulkanRHI->GetLogicalDevice(), &bufferDeviceAddressInfo);
	assert(deviceAddress != 0);
	return deviceAddress;
}

uint32 RHIBuffer::GetFrameCount() const
{
	return mIsMultiple ? GRenderState->GetFrameCount() : 1;
}

uint32 RHIBuffer::GetFrameIndex() const
{
	return mIsMultiple ? GRenderState->GetFrameIndex() : 0;
}
